import os
from typing import AsyncIterator, Dict, Any
import asyncio

from dotenv import load_dotenv
from langchain_openai import OpenAIEmbeddings
# 使用新的方式导入
from langchain.chains import ConversationalRetrievalChain
from langchain_openai import ChatOpenAI
from langchain.prompts import PromptTemplate
from langchain_chroma import Chroma


from chat_history import MultiUserConversationManager
from pymilvus import MilvusClient

# 加载环境变量
load_dotenv()
API_KEY = os.getenv("OPENAI_API_KEY")
BASE_URL = os.getenv("OPENAI_PROXY")

class LangChainLLMModule:
    def __init__(self):
        # 初始化Embedding模型（用于文本向量化）
        self.embeddings = OpenAIEmbeddings(model="text-embedding-3-large",api_key=API_KEY,base_url=BASE_URL)

        # 初始化大模型（以GPT-4为例）
        self.llm = ChatOpenAI(model="gpt-4o-mini-search-preview",api_key=API_KEY,base_url=BASE_URL,streaming=True)

        # 初始化多用户对话记忆管理器（使用Redis）
        self.conversation_manager = MultiUserConversationManager()

        # 定义RAG prompt模板（融合领域知识和对话历史）
        self.prompt_template = """
            你是一位专业的汽车推荐顾问，具备丰富的汽车行业知识和用户沟通经验。你的任务是根据用户的提问、历史对话以及知识库信息，提供专业、准确且易于理解的汽车推荐或解答。
            
            ### 一、信息调用规则
                1. 优先基于以下**历史对话**理解用户需求上下文，明确用户过往关注的预算、用途、偏好等关键信息：
                {chat_history}
            
                2. 严格依据以下**汽车知识库**内容输出数据，所有参数（如价格、油耗）、车型特点均需源于此，禁止编造：
                    {context}
            
          
            
            ### 二、核心任务：回答用户问题
                {question}
            
            ### 三、分场景回答规范
            #### 场景0：用户未提及汽车相关需求（如仅说姓名、打招呼、无关闲聊）
            1. 回应原则：先友好承接用户信息，再主动引导明确汽车需求，不主动推荐任何车型。
            2. 示例话术：
               - 若用户说姓名：“您好[用户姓名]！目前还没了解到您的汽车相关需求，请问您是想咨询具体车型，还是需要根据预算、用途推荐车型呢？”
               - 若用户打招呼（如“你好”）：“您好！我是专业的汽车推荐顾问，请问您有具体车型想了解，或需要根据预算、用途推荐家用/通勤车型吗？”
               
            #### 场景1：用户询问具体车型（如“丰田凯美瑞怎么样”）
            1. 基础信息：按“价格区间（指导价）、动力类型（燃油/混动/纯电）、百公里能耗（油耗X L/电耗X kWh）、车身尺寸（长×宽×高，单位mm）、轴距（mm）”的固定顺序呈现核心参数，参数需完整且来自知识库。
            2. 客观分析：分“优势”“不足”两点，每点用1-2个具体理由支撑（如优势：“混动版油耗低，市区仅5.3L/100km，适合日常通勤”；不足：“后排中间地台凸起较高，满载3人时中间乘客舒适性一般”）。
            3. 限制：仅围绕该车型展开，不额外推荐其他车型。
            
            #### 场景2：用户请求购车建议（如“20万预算买家用轿车”）
            1. 需求匹配：先明确用户提及的核心条件（预算、用途、燃油/新能源偏好等），若未提及用途，默认按“家用”场景推荐。
                        若用户未提及任何预算、用途等汽车需求（仅个人信息/闲聊），不进入本场景，转至“场景0”处理。
            2. 车型筛选：
                - 预算明确时：优先推荐价格在“预算±10%”范围内的车型；预算模糊时（如“十几万”），按“10-15万”等区间锚定；**无预算时不主动推荐车型**，转至场景0引导用户明确需求。
                - 知识库中fuel_type只要包含了电的都是新能源车。
                - 品牌与类型：至少推荐5款不同品牌车型，涵盖燃油车、新能源车（混动/纯电），同一品牌不重复推荐车型，不能推荐4或者5个同一品牌的汽车。
                - 价格校验：若用户未提预算，禁止推荐单价超25万的车型；若后续用户明确预算（如“20万左右”），严格排除超预算10%（即超22万）的车型。
                - 类型匹配：若用户未提车型类型（轿车/SUV），需均衡推荐（如5款车型中至少2款轿车、2款SUV），避免单一类型过度推荐。
            3. 推荐格式（每款车型独立成段）：
                - 车型名称：XX品牌 XX车型（如“本田 雅阁 2024款 锐·混动 锐领版”）
                - 关键参数：价格、汽车类型、动力类型、驱动方式、百公里加速（仅列此5项，信息源于知识库，如何知识库中没有，则可以不显示该属性）
                - 推荐理由：结合用户需求，说明“为何适合”（如“20万预算内，混动版油耗低且空间大，适合家庭长途出行”）
            
            #### 场景3：非汽车相关问题（如“今天天气如何”）
            1. 回应原则：友好告知“该问题与汽车无关",但可以适当对该问题扩展，此时回答无需包含汽车。
            2. 补充：若用户误发非汽车问题，可主动询问“是否有汽车相关问题需要咨询”，引导回归核心需求。
            
            4. **输出风格要求**：
               - 语言专业但通俗易懂
               - 使用清晰的条目结构（汽车使用1.编号、汽车参数使用项目符号②编号）
               - 避免冗长段落，突出重点信息
            """

        self.prompt = PromptTemplate(
            template=self.prompt_template,
            input_variables=["context", "question", "chat_history"]
        )

        self.vector_store = Chroma(
            persist_directory="../../asset/chroma_db",  # 本地存储路径
            embedding_function=self.embeddings
        )

        # 初始化RetrievalQA链（RAG核心逻辑）
        self.qa_chain = ConversationalRetrievalChain.from_llm(
            llm=self.llm,
            retriever=self.vector_store.as_retriever(search_kwargs={
                "k": 10,
            }),
            combine_docs_chain_kwargs={"prompt": self.prompt},
            return_source_documents=True
        )

    def generate_response(self, user_query, user_id=None):
        """生成包含RAG和对话历史的回答"""
        if user_id:
            # 获取用户的聊天历史记录
            chat_history = self.conversation_manager.get_message_history(user_id).messages

            # 转换为元组格式的对话历史 [(human_message, ai_message), ...]
            formatted_history = []
            for i in range(0, len(chat_history) - 1, 2):
                if i + 1 < len(chat_history):
                    human_msg = chat_history[i].content if hasattr(chat_history[i], 'content') else str(chat_history[i])
                    ai_msg = chat_history[i + 1].content if hasattr(chat_history[i + 1], 'content') else str(
                        chat_history[i + 1])
                    formatted_history.append((human_msg, ai_msg))

            # 调用 QA 链并传入聊天历史
            response = self.qa_chain.invoke({
                "question": user_query,
                "chat_history": formatted_history
            })

            # 保存用户和AI的对话记录
            self.conversation_manager.add_user_message(user_id, user_query)
            result_text = response.get("answer", str(response)) if isinstance(response, dict) else str(response)
            self.conversation_manager.add_ai_message(user_id, result_text)

            return result_text
        else:
            # 如果没有提供user_id，直接查询
            response = self.qa_chain.invoke({
                "question": user_query,
                "chat_history": []
            })
            if isinstance(response, dict):
                return response.get("answer", str(response))
            else:
                return str(response)

    async def generate_stream_response(self, user_query: str, user_id: str = None) -> AsyncIterator[str]:
        try:
            # 先保存用户消息（若有user_id）
            if user_id:
                self.conversation_manager.add_user_message(user_id, user_query)
                # 基于role获取历史，而非索引
                chat_history = self.conversation_manager.get_message_history(user_id).messages
                print(chat_history)
                formatted_history = []
                human_msg = None
                for msg in chat_history:
                    if not hasattr(msg, 'content') or not hasattr(msg, 'role'):
                        continue  # 跳过无效消息
                    if msg.role == 'user':
                        human_msg = msg.content
                    elif msg.role == 'assistant' and human_msg is not None:
                        formatted_history.append((human_msg, msg.content))
                        human_msg = None  # 配对后重置，避免重复
                # 若最后一条是未配对的用户消息，可单独处理（如作为上下文）
                # ...

            # 流式生成响应
            full_response = ""
            async for chunk in self.qa_chain.astream({"question": user_query, "chat_history": formatted_history}):
                if "answer" in chunk and isinstance(chunk["answer"], str):
                    content = chunk["answer"]
                    full_response += content
                    yield content  # 可根据需要调整粒度

            # 保存AI响应
            if user_id:
                self.conversation_manager.add_ai_message(user_id, full_response)

        except Exception as e:
            # 捕获异常并返回错误提示
            yield f"处理失败：{str(e)}"
            # 记录日志
            self.logger.error(f"Stream response error: {e}")

    async def generate_stream_response1(self, user_query: str, user_id: str = None) -> AsyncIterator[str]:
        """
        流式生成包含RAG和对话历史的回答

        Args:
            user_query: 用户的查询问题
            user_id: 用户标识符

        Yields:
            逐步生成的响应文本片段
        """
        if user_id:
            # 获取用户的聊天历史记录
            chat_history = self.conversation_manager.get_message_history(user_id).messages

            # 转换为元组格式的对话历史 [(human_message, ai_message), ...]
            formatted_history = []
            for i in range(0, len(chat_history) - 1, 2):
                if i + 1 < len(chat_history):
                    human_msg = chat_history[i].content if hasattr(chat_history[i], 'content') else str(chat_history[i])
                    ai_msg = chat_history[i + 1].content if hasattr(chat_history[i + 1], 'content') else str(
                        chat_history[i + 1])
                    formatted_history.append((human_msg, ai_msg))
        else:
            formatted_history = []

        # 构建输入
        inputs = {
            "question": user_query,
            "chat_history": formatted_history
        }

        # 收集完整的响应
        full_response = ""

        # 使用异步流式调用
        async for chunk in self.qa_chain.astream(inputs):
            if "answer" in chunk:
                content = chunk["answer"]
                full_response += content
                yield chunk["answer"]

        # 如果有user_id，则保存完整的对话记录
        if user_id:
            # 直接保存用户消息和AI响应，避免重复调用generate_response
            self.conversation_manager.add_user_message(user_id, user_query)
            self.conversation_manager.add_ai_message(user_id, full_response)

    def check_vector_store_quality(self,type):
        """检查现有向量存储质量"""
        try:
            # 获取一些样本进行分析
            sample_docs = self.vector_store.similarity_search(type, k=5)
            return sample_docs
        except Exception as e:
            print(f"检查向量存储时发生错误: {e}")
            return []
